from enum import Enum
import pandas as pd
from gofishing.base import TradeType
from gofishing.base import Uninstantiable, BaseBacktest, BaseRealtimeTQ, Log
from gofishing.base import generateID
from gofishing.commodity import CommodityBase


class OrderStatus(Enum):

    ALIVE = 'ALIVE'
    FINISHED = 'FINISHED'
    ERROR = 'ERROR'


class OrderInfo:

    def __init__(self,
                 orderID: str,
                 contractCode: str,
                 tradeType: TradeType,
                 status: OrderStatus,
                 limitPrice: float,
                 tradePrice: float,
                 volume: int,
                 volumeLeft: int,
                 insertTime: pd.Timestamp,
                 tradeTime: pd.Timestamp,
                 commission: float) -> None:
        self.orderID = orderID
        self.contractCode = contractCode
        self.tradeType = tradeType
        self.status = status
        self.limitPrice = limitPrice
        self.tradePrice = tradePrice
        self.volume = volume
        self.volumeLeft = volumeLeft
        self.insertTime = insertTime
        self.tradeTime = tradeTime
        self.commission = commission


class AccountValues:

    def __init__(self,
                 staticValue: float,
                 frozenMargin: float,
                 frozenCommission: float,
                 floatingProfit: float,
                 value: float,
                 margin: float,
                 available: float,
                 riskRatio: float) -> None:

        self.staticValue = staticValue
        self.frozenMargin = frozenMargin
        self.frozenCommission = frozenCommission

        self.floatingProfit = floatingProfit
        self.value = value
        self.margin = margin
        self.available = available

        self.riskRatio = riskRatio


class BacktestOrder:

    orders: dict[str, 'BacktestOrder'] = {}
    aliveOrderIDs: list[str] = []

    def __init__(self, contractCode: str, tradeType: TradeType,
                 limitPrice: float, volume: int) -> None:
        self.contractCode = contractCode
        self.tradeType = tradeType
        self.status = OrderStatus.ALIVE
        self.limitPrice = limitPrice
        self.tradePrice = 0.0
        self.volume = volume
        self.volumeLeft = self.volume
        self.insertTime = BaseBacktest.now()
        self.tradeTime = pd.Timestamp('')
        self.commission = 0.0
        self.orderID = generateID()
        BacktestOrder.orders[self.orderID] = self

    def _updateAliveOrderIDs() -> None:
        BacktestOrder.aliveOrderIDs = []
        orders = BacktestOrder.orders
        for orderID in orders:
            if orders[orderID].status == OrderStatus.ALIVE:
                BacktestOrder.aliveOrderIDs.append(orderID)

    def update() -> None:
        BacktestOrder._updateAliveOrderIDs()


OrderB = BacktestOrder


class BrokerBase(Uninstantiable):

    def update() -> None:
        raise NotImplementedError

    def order(commodityID: str, tradeType: TradeType,
              limitPrice: float, volume: int) -> str:
        raise NotImplementedError

    def getOrderInfo(orderID: str) -> OrderInfo:
        raise NotImplementedError

    def getAccountValues() -> AccountValues:
        raise NotImplementedError

    def getHolding() -> dict:
        raise NotImplementedError


class BrokerBacktest(BrokerBase):

    holding: dict[str, dict] = {}
    staticValue: float = 50000.0  # 静态权益, 只有在买/卖发生时才更新
    frozenMargin: float = 0.0  # 冻结保证金
    frozenCommission: float = 0.0  # 冻结手续费

    # 这几个需由上面的信息计算得到
    floatingProfit: float = 0.0  # 浮动盈亏
    value: float = 50000.0  # 客户权益
    margin: float = 0.0  # 保证金
    available: float = 50000.0  # 可用资金

    riskRatio: float = 0.0  # 风险度

    def _addHoldingInfo(contractCode: str) -> None:
        info = {'long': 0,
                'short': 0,
                'longFrozen': 0,  # 委托平仓时冻结的仓位
                'shortFrozen': 0,  # 委托平仓时冻结的仓位
                'longOpenPrice': 0.0,
                'shortOpenPrice': 0.0}
        BrokerBacktest.holding[contractCode] = info

    def _cleanEmptyHolding() -> None:
        holding = BrokerBacktest.holding
        contractEmpty = []
        aliveContracts = []
        for orderID in OrderB.aliveOrderIDs:
            aliveContracts.append(OrderB.orders[orderID].contractCode)
        for contractCode in holding:
            holdingInfo = holding[contractCode]
            if not holdingInfo['long'] and not holdingInfo['short'] and \
                    contractCode not in aliveContracts:
                contractEmpty.append(contractCode)
        for contractCode in contractEmpty:
            holding.pop(contractCode)

    def _holdingFreeze(contractCode: str,
                       longFrozenDiff: int, shortFrozenDiff: int) -> None:
        holdingInfo = BrokerBacktest.holding[contractCode]
        longFrozen = holdingInfo['longFrozen'] + longFrozenDiff
        shortFrozen = holdingInfo['shortFrozen'] + shortFrozenDiff
        holdingInfo['longFrozen'] = longFrozen
        holdingInfo['shortFrozen'] = shortFrozen

    def _holdingChange(contractCode: str,
                       tradePrice: float,
                       longAdd=0, longCut=0,
                       shortAdd=0, shortCut=0) -> None:
        holdingInfo = BrokerBacktest.holding[contractCode]
        longPre = holdingInfo['long']
        longOpenPricePre = holdingInfo['longOpenPrice']
        shortPre = holdingInfo['short']
        shortOpenPricePre = holdingInfo['shortOpenPrice']
        if not longAdd and not longCut and not shortAdd and not shortCut:
            raise Exception('must set a position change')
        elif not longCut and not shortAdd and not shortCut:
            long = longPre + longAdd
            longOpenPrice = \
                (longPre * longOpenPricePre + longAdd * tradePrice) / long
            holdingInfo['long'] = long
            holdingInfo['longOpenPrice'] = longOpenPrice
        elif not longAdd and not shortAdd and not shortCut:
            long = longPre - longCut
            holdingInfo['long'] = long
            if not long:
                holdingInfo['longOpenPrice'] = 0.0
        elif not longAdd and not longCut and not shortCut:
            short = shortPre + shortAdd
            shortOpenPrice = \
                (shortPre * shortOpenPricePre + shortAdd * tradePrice) / short
            holdingInfo['short'] = short
            holdingInfo['shortOpenPrice'] = shortOpenPrice
        elif not longAdd and not longCut and not shortAdd:
            short = shortPre - shortCut
            holdingInfo['short'] = short
            if not short:
                holdingInfo['shortOpenPrice'] = 0.0
        else:
            raise Exception('just one kind of position change is ok')

    def getFloatingProfit(contractCode: str) -> float:
        holdingInfo = BrokerBacktest.holding[contractCode]
        commodity = CommodityBase.findCommodity(contractCode=contractCode)
        long = holdingInfo['long']
        longOpenPrice = holdingInfo['longOpenPrice']
        short = holdingInfo['short']
        shortOpenPrice = holdingInfo['shortOpenPrice']
        mult = commodity.mult
        close = commodity.close[-1]
        floatingProfitLong = long * mult * (close - longOpenPrice)
        floatingProfitShort = short * mult * (shortOpenPrice - close)
        return floatingProfitLong + floatingProfitShort

    def getMargin(contractCode: str) -> float:
        holdingInfo = BrokerBacktest.holding[contractCode]
        commodity = CommodityBase.findCommodity(contractCode=contractCode)
        long = holdingInfo['long']
        short = holdingInfo['short']
        mult = commodity.mult
        close = commodity.close[-1]
        marginRate = commodity.marginRate
        return long * (close * mult * marginRate) + \
            short * (close * mult * marginRate)

    def _updateOrdering() -> None:
        orders = OrderB.orders
        aliveOrderIDs = OrderB.aliveOrderIDs
        for orderID in aliveOrderIDs:
            orderB = orders[orderID]
            tradeType = orderB.tradeType
            if tradeType in [TradeType.LONGOPEN, TradeType.SHORTOPEN]:
                BrokerBacktest._finishOpenOrder(orderB=orderB)
            elif tradeType in [TradeType.LONGCLOSE, TradeType.SHORTCLOSE]:
                BrokerBacktest._finishCloseOrder(orderB=orderB)

    def _updateValues() -> None:
        BrokerBacktest.floatingProfit = 0.0
        BrokerBacktest.margin = 0.0
        for contractCode in BrokerBacktest.holding:
            BrokerBacktest.floatingProfit += \
                BrokerBacktest.getFloatingProfit(contractCode)
            BrokerBacktest.margin += BrokerBacktest.getMargin(contractCode)
        BrokerBacktest.value = \
            BrokerBacktest.staticValue + BrokerBacktest.floatingProfit
        BrokerBacktest.available = \
            BrokerBacktest.value - BrokerBacktest.margin - \
            BrokerBacktest.frozenMargin - BrokerBacktest.frozenCommission
        BrokerBacktest.riskRatio = BrokerBacktest.margin / BrokerBacktest.value

    def getAccountValues() -> AccountValues:
        BrokerB = BrokerBacktest
        BrokerB._updateValues()
        return AccountValues(staticValue=BrokerB.staticValue,
                             frozenMargin=BrokerB.frozenMargin,
                             frozenCommission=BrokerB.frozenCommission,
                             floatingProfit=BrokerB.floatingProfit,
                             value=BrokerB.value,
                             margin=BrokerB.margin,
                             available=BrokerB.available,
                             riskRatio=BrokerB.riskRatio)

    def getOrderInfo(orderID: str) -> OrderInfo:
        orderB: OrderB = OrderB.orders[orderID]
        return OrderInfo(orderID=orderID,
                         contractCode=orderB.contractCode,
                         tradeType=orderB.tradeType,
                         status=orderB.status,
                         limitPrice=orderB.limitPrice,
                         tradePrice=orderB.tradePrice,
                         volume=orderB.volume,
                         volumeLeft=orderB.volumeLeft,
                         insertTime=orderB.insertTime,
                         tradeTime=orderB.tradeTime,
                         commission=orderB.commission)

    def getHolding() -> dict:
        holding = {}
        for contractCode in BrokerBacktest.holding:
            holdingInfo: dict = BrokerBacktest.holding[contractCode].copy()
            holdingInfo.pop('longFrozen')
            holdingInfo.pop('shortFrozen')
            holding[contractCode] = holdingInfo
        return holding

    def update() -> None:
        OrderB.update()
        BrokerBacktest._updateOrdering()
        OrderB.update()
        BrokerBacktest._updateValues()
        BrokerBacktest._cleanEmptyHolding()

    def order(contractCode: str, tradeType: TradeType,
              limitPrice: float, volume: int) -> str:
        orderB = OrderB(contractCode=contractCode, tradeType=tradeType,
                        limitPrice=limitPrice, volume=volume)
        if contractCode not in BrokerBacktest.holding:
            BrokerBacktest._addHoldingInfo(contractCode=contractCode)

        # TODO 如果是收盘时间，就直接 order error
        if tradeType in [TradeType.LONGOPEN, TradeType.SHORTOPEN]:
            BrokerBacktest._open(orderB=orderB)
        elif tradeType in [TradeType.LONGCLOSE, TradeType.SHORTCLOSE]:
            BrokerBacktest._close(orderB=orderB)
        return orderB.orderID

    def _open(orderB: OrderB) -> None:
        commodity = CommodityBase.findCommodity(orderB.contractCode)
        limitPrice = orderB.limitPrice
        volume = orderB.volume
        frozenMargin = volume * \
            (limitPrice * commodity.mult * commodity.marginRate +
             commodity.openCommission)
        frozenCommission = volume * commodity.openCommission
        if BaseBacktest.isCloseOfDay():
            orderB.status = OrderStatus.ERROR
        elif frozenMargin + frozenCommission <= BrokerBacktest.available:
            BrokerBacktest.frozenMargin += frozenMargin
            BrokerBacktest.frozenCommission += frozenCommission
            BrokerBacktest._updateValues()
        else:
            Log.warning('[orderID %s] available not enough to open order' %
                        (orderB.orderID))

    def _finishOpenOrder(orderB: OrderB) -> None:
        now = BaseBacktest.now()
        contractCode = orderB.contractCode
        commodity = CommodityBase.findCommodity(contractCode=contractCode)
        tradePrice = orderB.limitPrice  # 模拟全部按委托价成交
        tradeVolume = orderB.volume  # 模拟全部按委托价成交
        reasonableRange = commodity.oneBarPV(start=orderB.insertTime,
                                             end=now)

        unFrozenMargin = tradeVolume * \
            (tradePrice * commodity.mult * commodity.marginRate +
                commodity.openCommission)
        commission = tradeVolume * commodity.openCommission

        if tradePrice > reasonableRange['low'][-1] and \
                tradePrice < reasonableRange['high'][-1] and \
                tradeVolume <= reasonableRange['volume'][-1]:

            tradeType = orderB.tradeType
            if tradeType == TradeType.LONGOPEN:
                BrokerBacktest._holdingChange(contractCode=contractCode,
                                              tradePrice=tradePrice,
                                              longAdd=tradeVolume)
            elif tradeType == TradeType.SHORTOPEN:
                BrokerBacktest._holdingChange(contractCode=contractCode,
                                              tradePrice=tradePrice,
                                              shortAdd=tradeVolume)

            BrokerBacktest.staticValue -= commission
            BrokerBacktest.frozenMargin -= unFrozenMargin
            BrokerBacktest.frozenCommission -= commission
            BrokerBacktest._updateValues()

            orderB.volumeLeft -= tradeVolume
            orderB.tradePrice = tradePrice
            orderB.tradeTime = now
            orderB.commission = commission
            orderB.status = OrderStatus.FINISHED

        elif BaseBacktest.isCloseOfDay():
            BrokerBacktest.frozenMargin -= unFrozenMargin
            BrokerBacktest.frozenCommission -= commission
            BrokerBacktest._updateValues()
            orderB.status = OrderStatus.ERROR

        else:
            # 未成交
            pass

    def _close(orderB: OrderB) -> None:
        contractCode = orderB.contractCode
        commodity = CommodityBase.findCommodity(contractCode=contractCode)
        volume = orderB.volume
        frozenCommission = volume * commodity.closeCommission
        holdingInfo = BrokerBacktest.holding[contractCode]
        tradeType = orderB.tradeType
        if BaseBacktest.isCloseOfDay():
            orderB.status = OrderStatus.ERROR
        elif tradeType == TradeType.LONGCLOSE:
            longFree = holdingInfo['long'] - holdingInfo['longFrozen']
            if longFree >= volume:
                BrokerBacktest._holdingFreeze(
                    contractCode=contractCode,
                    longFrozenDiff=volume,
                    shortFrozenDiff=0)
                BrokerBacktest.frozenCommission += frozenCommission
                BrokerBacktest._updateValues()
            else:
                Log.warning('[orderID %s] longFree not enough for closing' %
                            (orderB.orderID))
        elif tradeType == TradeType.SHORTCLOSE:
            shortFree = holdingInfo['short'] - holdingInfo['shortFrozen']
            if shortFree >= volume:
                BrokerBacktest._holdingFreeze(
                    contractCode=contractCode,
                    longFrozenDiff=0,
                    shortFrozenDiff=volume)
                BrokerBacktest.frozenCommission += frozenCommission
                BrokerBacktest._updateValues()
            else:
                Log.warning('[orderID %s] shortFree not enough for closing' %
                            (orderB.orderID))

    def _finishCloseOrder(orderB: OrderB) -> None:
        now = BaseBacktest.now()
        contractCode = orderB.contractCode
        commodity = CommodityBase.findCommodity(contractCode=contractCode)
        tradeType = orderB.tradeType
        tradePrice = orderB.limitPrice  # 模拟全部按委托价成交
        tradeVolume = orderB.volume  # 模拟全部按委托价成交
        reasonableRange = commodity.oneBarPV(start=orderB.insertTime,
                                             end=now)

        commission = tradeVolume * commodity.closeCommission

        if tradePrice > reasonableRange['low'][-1] and \
                tradePrice < reasonableRange['high'][-1] and \
                tradeVolume <= reasonableRange['volume'][-1]:
            mult = commodity.mult
            holdingInfo = BrokerBacktest.holding[contractCode]

            if tradeType == TradeType.LONGCLOSE:
                longOpenPrice = holdingInfo['longOpenPrice']
                profit = tradeVolume * (tradePrice - longOpenPrice) * mult
                BrokerBacktest.staticValue += profit - commission
                BrokerBacktest._holdingChange(contractCode=contractCode,
                                              tradePrice=tradePrice,
                                              longCut=tradeVolume)
                BrokerBacktest._holdingFreeze(contractCode=contractCode,
                                              longFrozenDiff=(-tradeVolume),
                                              shortFrozenDiff=0)
            elif tradeType == TradeType.SHORTCLOSE:
                shortOpenPrice = holdingInfo['shortOpenPrice']
                profit = tradeVolume * (shortOpenPrice - tradePrice) * mult
                BrokerBacktest.staticValue += profit - commission
                BrokerBacktest._holdingChange(contractCode=contractCode,
                                              tradePrice=tradePrice,
                                              shortCut=tradeVolume)
                BrokerBacktest._holdingFreeze(contractCode=contractCode,
                                              longFrozenDiff=0,
                                              shortFrozenDiff=(-tradeVolume))
            BrokerBacktest.frozenCommission -= commission
            BrokerBacktest._updateValues()

            orderB.volumeLeft -= tradeVolume
            orderB.tradePrice = tradePrice
            orderB.tradeTime = BaseBacktest.now()
            orderB.commission = commission
            orderB.status = OrderStatus.FINISHED

        elif BaseBacktest.isCloseOfDay():
            if tradeType == TradeType.LONGCLOSE:
                BrokerBacktest._holdingFreeze(contractCode=contractCode,
                                              longFrozenDiff=(-tradeVolume),
                                              shortFrozenDiff=0)
            elif tradeType == TradeType.SHORTCLOSE:
                BrokerBacktest._holdingFreeze(contractCode=contractCode,
                                              longFrozenDiff=0,
                                              shortFrozenDiff=(-tradeVolume))
            BrokerBacktest.frozenCommission -= commission
            BrokerBacktest._updateValues()
            orderB.status = OrderStatus.ERROR

        else:
            # 未成交
            pass


class BrokerRealtimeTQ(BrokerBase):

    def update() -> None:
        pass

    def order(contractCode: str, tradeType: TradeType,
              limitPrice: float, volume: int) -> str:
        tqApi = BaseRealtimeTQ.tqApi
        BrokerTQ = BrokerRealtimeTQ
        direction, offset = \
            BrokerTQ._tradeTypeToDirectionOffset(tradeType)
        order = tqApi.insert_order(symbol=contractCode,
                                   direction=direction,
                                   offset=offset,
                                   volume=volume,
                                   limit_price=limitPrice)
        return order.order_id

    def getOrderInfo(orderID: str) -> OrderInfo:
        tqApi = BaseRealtimeTQ.tqApi
        BrokerTQ = BrokerRealtimeTQ
        order = tqApi.get_order(orderID)
        return OrderInfo(
            orderID=orderID,
            contractCode=order.instrument_id,
            tradeType=BrokerTQ._directionOffsetToTradeType(
                direction=order.direction, offset=order.offset),
            status=BrokerTQ._toOrderStatus(orderTQ=order),
            limitPrice=order.limit_price,
            tradePrice=order.trade_price,
            volume=order.volume_orign,
            volumeLeft=order.volume_left,
            insertTime=BrokerTQ._datetimeToTimestamp(order.insert_date_time),
            tradeTime=pd.Timestamp(0),  # TODO 待完善
            commission=0.0)

    def getAccountValues() -> AccountValues:
        tqApi = BaseRealtimeTQ.tqApi
        account = tqApi.get_account()
        return AccountValues(
            staticValue=account.static_balance,  # 注意，和回测有区别
            frozenMargin=account.frozen_margin,
            frozenCommission=account.frozen_commission,
            floatingProfit=account.float_profit,
            value=account.balance,
            margin=account.margin,
            available=account.available,
            riskRatio=account.risk_ratio)

    def getHolding() -> dict:
        tqApi = BaseRealtimeTQ.tqApi
        positions = tqApi.get_position()
        holding = {}
        for contractCode in positions:
            position = positions[contractCode]
            holdingInfo = {'long': position.pos_long,
                           'short': position.pos_short,
                           'longOpenPrice': position.open_price_long,
                           'shortOpenPrice': position.open_price_short}
            holding[contractCode] = holdingInfo
        return holding

    def _tradeTypeToDirectionOffset(tradeType: TradeType) -> tuple[str, str]:
        if tradeType == TradeType.LONGOPEN:
            return ('BUY', 'OPEN')
        elif tradeType == TradeType.LONGCLOSE:
            return ('SELL', 'CLOSE')
        elif tradeType == TradeType.SHORTOPEN:
            return ('SELL', 'OPEN')
        elif tradeType == TradeType.SHORTCLOSE:
            return ('BUY', 'CLOSE')
        else:
            raise Exception('tradeType error')

    def _directionOffsetToTradeType(direction: str, offset: str) -> TradeType:
        if direction == 'BUY':
            if offset == 'OPEN':
                return TradeType.LONGOPEN
            elif offset == 'CLOSE':
                return TradeType.SHORTCLOSE
        elif direction == 'SELL':
            if offset == 'OPEN':
                return TradeType.SHORTOPEN
            elif offset == 'CLOSE':
                return TradeType.LONGCLOSE
        else:
            raise Exception('directionOffset error')

    def _datetimeToTimestamp(ns: int) -> pd.Timestamp:
        return pd.Timestamp(ns)

    def _toOrderStatus(orderTQ) -> OrderStatus:
        if orderTQ.is_dead or orderTQ.is_error:
            return OrderStatus.ERROR
        elif orderTQ.status == 'ALIVE':
            return OrderStatus.ALIVE
        elif orderTQ.status == 'FINISHED':
            return OrderStatus.FINISHED
        else:
            raise Exception('status error')


Broker = BrokerBacktest


def setBrokerRealtime() -> None:
    global Broker
    Broker = BrokerRealtimeTQ
